package org.glandais.profiler;

import java.io.IOException;
import java.util.HashSet;
import java.util.Set;

import javassist.CannotCompileException;
import javassist.ClassPath;
import javassist.ClassPool;
import javassist.CtBehavior;
import javassist.CtClass;
import javassist.CtMethod;
import javassist.LoaderClassPath;
import javassist.Modifier;
import javassist.NotFoundException;
import javassist.bytecode.MethodInfo;

public class ClassFileTransformer implements java.lang.instrument.ClassFileTransformer {

	private static final String ORG_GLANDAIS_PROFILER = "org.glandais.profiler"; //$NON-NLS-1$
	private static final String JAVA_LANG_THROWABLE = "java.lang.Throwable"; //$NON-NLS-1$
	private ClassPool classPool = ClassPool.getDefault();
	private Set<ClassLoader> loaders = new HashSet<ClassLoader>();
	private CtClass ctThrowable = classPool.get(JAVA_LANG_THROWABLE);

	public static boolean doTransformClasses = true;

	class BehaviorInfo {

		private static final String STAT = "STAT"; //$NON-NLS-1$
		private static final String INIT = "INIT"; //$NON-NLS-1$
		private static final String DOLLAR_0 = "$0"; //$NON-NLS-1$
		private static final String METH = "METH"; //$NON-NLS-1$
		private static final String DOT = "."; //$NON-NLS-1$
		private static final String STATIC = "static "; //$NON-NLS-1$
		private static final String ACCOLADES = "{} "; //$NON-NLS-1$
		private static final String SPACE = " "; //$NON-NLS-1$
		private static final String NEW = "new "; //$NON-NLS-1$
		private static final String CONS = "CONS"; //$NON-NLS-1$
		private static final String UNKN = "UNKN"; //$NON-NLS-1$
		private static final String NULL = "null"; //$NON-NLS-1$
		String clazz = null;
		String methodSignature = null;
		String sig = null;
		String sthis = NULL;
		String type = UNKN;

		public BehaviorInfo(CtBehavior method) {
			super();
			MethodInfo methodInfo2 = method.getMethodInfo2();
			clazz = method.getDeclaringClass().getName();
			methodSignature = method.getSignature();
			if (methodInfo2.isConstructor()) {
				type = CONS;
				sig = NEW + clazz + SPACE + methodSignature;
			} else if (methodInfo2.isStaticInitializer()) {
				type = INIT;
				sig = clazz + ACCOLADES + methodSignature;
			} else if (Modifier.isStatic(method.getModifiers())) {
				type = STAT;
				sig = STATIC + clazz + DOT + method.getName() + SPACE + methodSignature;
			} else if (method instanceof CtMethod) {
				type = METH;
				sthis = DOLLAR_0;
				sig = clazz + DOT + method.getName() + SPACE + methodSignature;
			}
		}

	}

	public ClassFileTransformer(String agentArgument) throws Exception {
		super();
	}

	public final byte[] loadAndTransform(ClassLoader loader, String className, Class<?> clazz,
			java.security.ProtectionDomain domain) throws IOException, CannotCompileException, NotFoundException {
		if (!loaders.contains(loader)) {
			addClassLoader(loader);
		}
		CtClass ctClass = classPool.get(className);
		return doClass(loader, className, clazz, ctClass);
	}

	public final boolean shouldTransform(String className) {
		if (doTransformClasses) {
			for (String excluded : Agent.configBean.getExcluded()) {
				if (className.replace('/', '.').startsWith(excluded.replace('/', '.'))) {
					if (Agent.configBean.isLogTraceInstrumentation()) {
						Agent.configBean.logTraceInstrumentation(Messages.getString("ClassFileTransformer.EXCLUDED") + className); //$NON-NLS-1$
					}
					return false;
				}
			}
			if (className.replace('/', '.').startsWith(ORG_GLANDAIS_PROFILER)) {
				if (Agent.configBean.isLogTraceInstrumentation()) {
					Agent.configBean.logTraceInstrumentation(Messages.getString("ClassFileTransformer.EXCLUDED") + className); //$NON-NLS-1$
				}
				return false;
			}
			if (Agent.configBean.getIncluded().size() == 0) {
				return true;
			} else {
				for (String included : Agent.configBean.getIncluded()) {
					if (className.replace('/', '.').startsWith(included.replace('/', '.'))) {
						return true;
					}
				}
			}
		}
		return false;
	}

	@Override
	public final byte[] transform(ClassLoader loader, String className, Class<?> clazz,
			java.security.ProtectionDomain domain, byte[] bytes) {
		TraceRecorder.totalTimeTL.get().startProcessing();
		if (shouldTransform(className)) {
			byte[] result = doClass(loader, className, clazz, bytes);
			TraceRecorder.totalTimeTL.get().stopProcessing();
			return result;
		}
		TraceRecorder.totalTimeTL.get().stopProcessing();
		return bytes;
	}

	private synchronized final void addClassLoader(ClassLoader loader) {
		if (!loaders.contains(loader)) {
			if (Agent.configBean.isLogTraceInstrumentation()) {
				Agent.configBean.logTraceInstrumentation(Messages.getString("ClassFileTransformer.ADDING_CLASS_LOADER") + loader); //$NON-NLS-1$
			}
			loaders.add(loader);
			ClassPath cp = new LoaderClassPath(loader);
			classPool.appendClassPath(cp);
		}
	}

	private final byte[] doClass(ClassLoader loader, String name, Class<?> clazz, byte[] b) {
		CtClass cl = null;
		try {
			cl = classPool.makeClass(new java.io.ByteArrayInputStream(b));
			return doClass(loader, name, clazz, cl);
		} catch (Throwable e) {
			Agent.configBean.logError(Messages.getString("ClassFileTransformer.COULD_NOT_INSTRUMENT") + name, e); //$NON-NLS-1$
		}
		return new byte[0];
	}

	private final byte[] doClass(ClassLoader loader, String name, Class<?> clazz, CtClass cl) {
		if (!loaders.contains(loader)) {
			addClassLoader(loader);
		}

		if (Agent.configBean.isLogVerbose()) {
			Agent.configBean.logVerbose(Messages.getString("ClassFileTransformer.TRANSFORMING") + name); //$NON-NLS-1$
		}

		byte[] b = new byte[0];
		try {
			if (cl.isInterface() == false) {
				CtBehavior[] methods = cl.getDeclaredBehaviors();
				for (int i = 0; i < methods.length; i++) {
					if (methods[i].isEmpty() == false) {
						BehaviorInfo behaviorInfo = new BehaviorInfo(methods[i]);
						doMethod(behaviorInfo, methods[i]);
					}
				}
				b = cl.toBytecode();
			}
		} catch (Throwable e) {
			Agent.configBean.logError(Messages.getString("ClassFileTransformer.COULD_NOT_INSTRUMENT") + name, e); //$NON-NLS-1$
		} finally {
			if (cl != null) {
				cl.detach();
			}
		}
		return b;
	}

	private final void doMethod(BehaviorInfo behaviorInfo, CtBehavior method) throws NotFoundException,
			CannotCompileException {
		method.insertBefore("org.glandais.profiler.TraceRecorder.enter( \"" + behaviorInfo.type + "\", " //$NON-NLS-1$ //$NON-NLS-2$
				+ behaviorInfo.sthis + " ,\"" + behaviorInfo.sig + "\" );"); //$NON-NLS-1$ //$NON-NLS-2$
		String exitCode = "org.glandais.profiler.TraceRecorder.exit( \"" + behaviorInfo.type + "\", " //$NON-NLS-1$ //$NON-NLS-2$
				+ behaviorInfo.sthis + " , \"" + behaviorInfo.sig + "\" );"; //$NON-NLS-1$ //$NON-NLS-2$
		method.addCatch("{ " + exitCode + " throw $e; }", ctThrowable); //$NON-NLS-1$ //$NON-NLS-2$
		method.insertAfter(exitCode);
	}
}
